package com.gframework.mybatis.dao.mybatis.plugins;

import java.io.Closeable;
import java.lang.reflect.Field;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;

import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.logging.jdbc.ConnectionLogger;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.BeanInitializationException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.config.ConfigurableBeanFactory;
import org.springframework.context.annotation.Scope;
import org.springframework.core.annotation.Order;
import org.springframework.transaction.interceptor.TransactionAspectSupport;

import com.gframework.mybatis.dao.IMybatisDAO;
import com.gframework.mybatis.dao.datasource.DataSourceException;
import com.gframework.mybatis.dao.datasource.DynamicDataSourceManager;

/**
 * mybatis的多数据操作拦截器.<br>
 * 
 * <pre>
 * 执行任何一个 <strong>实现了{@link IMybatisDAO}接口的</strong> DAO bean方法前，
 * 修改数据源可以修改修改本次dao方法执行时操作的数据源（只读）——优先级较高
 * </pre>
 * <strong>使用安全性问题：【重要】</strong>
 * <p>
 * 请务必保证在调用本类方法后立刻去执行dao方法，这样就不会出现问题。
 * </p>
 * 
 * @since 1.0.0
 * @author Ghwolf
 */
@Order(10)
@Scope(ConfigurableBeanFactory.SCOPE_SINGLETON)
@Aspect
public class DataSourceHelper implements InitializingBean {

	private static final Logger logger = LoggerFactory.getLogger(DataSourceHelper.class);

	/**
	 * 本类唯一对象应用
	 */
	private static DataSourceHelper instance ;
	/**
	 * 多数据源获取和操作对象
	 */
	@Autowired
	private DynamicDataSourceManager dataSource;

	/**
	 * 当前使用的数据源key值.<br>
	 * 获取为null表示使用主数据源
	 */
	private final ThreadLocal<ConnectionCache> currentDataSource = new ThreadLocal<>();

	/**
	 * BoundSql类成员属性sql类型对象
	 */
	private Field sqlField;
	
	/**
	 * BaseJdbcLogger类成员属性connection对象
	 */
	private Field connectionField;
	
	/**
	 * mybatis插件对象
	 */
	private DataSourceHelperInterceptor dataSourceHelperInterceptor = new DataSourceHelperInterceptor();
	
	/**
	 * 多数据源数据库连接对象缓存类.
	 * 只有非主数据源才会使用到此类
	 * @author Ghwolf
	 *
	 */
	class ConnectionCache implements Closeable {
		/**
		 * 数据源key
		 */
		String key;
		/**
		 * 数据库连接对象
		 */
		Connection conn;
		/**
		 * 是否不允许修改，在执行{@link DataSourceHelper#doQuery(String, Supplier)}等类似方法过程中，是不允许在对数据源进行修改的，这个时候可以将其设置为true
		 */
		final boolean inQueryCanNotChange ;
		
		public ConnectionCache(boolean inQueryCanNotChange){
			this.inQueryCanNotChange = inQueryCanNotChange;
		}
		
		public boolean isInQueryCanNotChange() {
			return this.inQueryCanNotChange;
		}

		public String getKey() {
			return this.key;
		}
		public void setKey(String key) {
			this.key = key;
		}
		/**
		 * 获取数据库连接对象
		 * @return 返回Connection类实例化对象
		 * @throws SQLException 如果获取失败，则抛出此异常
		 */
		public Connection getConn() throws SQLException {
			if (this.conn == null) {
				Connection conn = DataSourceHelper.this.dataSource.openConnection(key);
				conn.setReadOnly(true);
				this.conn = conn;
			}
			return this.conn;
		}
		/**
		 * 使用此方法进行数据源的关闭，你不用再关注数据源对象是否存在的问题。也不必关系异常问题.
		 * 此方法会捕捉所有异常，同事如果数据源对象不存在，则不会做任何事情。
		 */
		@Override
		public void close() {
			try {
				if (this.conn != null &&  !this.conn.isClosed()) {
					this.conn.close();
				}
			} catch (SQLException e) {
				logger.error("关闭数据与出错：{}",this.key,e);
			}
		}

	}
	
	public DataSourceHelper() {
		try {
			this.sqlField = BoundSql.class.getDeclaredField("sql");
			this.connectionField = ConnectionLogger.class.getDeclaredField("connection");
		} catch (NoSuchFieldException e) {
			throw new BeanInitializationException("BoundSql类没有sql属性或ConnectionLogger类灭有connection属性，请检查mybatis版本是否符合要求！",e);
		} catch (SecurityException e) {
			throw new BeanInitializationException("存在安全管理器，无法反射获取类属性！",e);
		}
		this.sqlField.setAccessible(true);
		this.connectionField.setAccessible(true);
	}
	
	
	/**
	 * 获取mybatis插件对象
	 * @return {@link Interceptor}接口子类对象
	 */
	public DataSourceHelperInterceptor getDataSourceHelperInterceptor() {
		return this.dataSourceHelperInterceptor;
	}

	/**
	 * 拦截所有实现了IMybatisDAO接口的方法，在操作前后做数据源修改和回收处理。
	 */
	@Around("execution(* com.gframework.mybatis.dao.IMybatisDAO+.*(..))")
	public Object mybatisDaoAround(ProceedingJoinPoint point) throws Throwable{
		
		ConnectionCache conn = this.currentDataSource.get();
		if (conn == null) {
			return point.proceed(point.getArgs());
		} else if (conn.getKey() == null) {
			this.currentDataSource.remove();
			return point.proceed(point.getArgs());
		}
		
		try {
			return point.proceed(point.getArgs());
		} finally {
			// 执行doQuery等方法时，不通过aop回收
			if (!conn.isInQueryCanNotChange()) {
				this.currentDataSource.remove();
				conn.close();
			}
		}
	}
	
	
	/**
	 * 修改之后所有sql操作的数据源【仅支持查询操作，无事务】.
	 * <p>
	 * 请保证调用此方法后立刻去执行mybatis的dao操作，这样才不会出现任何问题。
	 * </p>
	 * 数据源的获取是通过{@link DynamicDataSourceManager}接口实现的。
	 * 
	 * @param dataSourceKey 要修改成为的数据源的id
	 *            {@link DynamicDataSourceManager}，如果不存在此数据源，则抛出异常
	 * @return 返回值可以让你知道你的设置是否覆盖了已有的配置。如果返回true表示这是一个全新的设置，没有覆盖任何设置，否则表示当前已经存在设置但是被你最新的设置给覆盖了。
	 * @throws DataSourceException 如果获取数据源发生异常，则抛出此异常
	 * @see DynamicDataSourceManager
	 */
	public static boolean setDataSource(String dataSourceKey) {
		return DataSourceHelper.instance.setDataSource0(dataSourceKey,false);
	}
	
	private static <T,U,R> R doQuery0(String dataSourceKey,BiFunction<T, U, R> function,T t,U u) {
		DataSourceHelper.instance.setDataSource0(dataSourceKey,true);
		try (ConnectionCache conn = DataSourceHelper.instance.currentDataSource.get()){
			if (conn == null) {
				throw new IllegalArgumentException("指定数据源不存在：" + dataSourceKey);
			}
			return function.apply(t, u);
		} finally {
			DataSourceHelper.instance.currentDataSource.remove();
		}
	}
	
	/**
	 * 修改数据源并进行一次相对应的操作.
	 * <p>数据源修改相关问题可以参考{@link #setDataSource(String)}方法。
	 * <p>不同于{@link #setDataSource(String)}的是，它可以避免因PageHelper和DataSourceHelper混合使用
	 * 导致开发者不知道先执行哪个造成的困惑。因为PageHelper和DataSourceHelper原理一致，都必须在执行完后立刻查询sql，如果中间出现异常，
	 * 则会出现问题。那么这个时候你就可以使用此方法。
	 * <p>同时此方法还有一点不同在于他的修改不会仅限于一次dao的查询，数据源修改的有效期会持续到整个方法结束。但是在此方法期间不允许在对数据源在进行修改，否则将会抛出异常！
	 * @param dataSourceKey
	 * @param run 要执行的操作
	 * @see #setDataSource(String)
	 */
	public static void doQuery(String dataSourceKey,Runnable run) {
		doQuery0(dataSourceKey, (a, b) -> {
			a.run();
			return null;
		} , run, null);
	}
	
	/**
	 * 修改数据源并进行一次相对应的操作.
	 * <p>数据源修改相关问题可以参考{@link #setDataSource(String)}方法。
	 * <p>不同于{@link #setDataSource(String)}的是，它可以避免因PageHelper和DataSourceHelper混合使用
	 * 导致开发者不知道先执行哪个造成的困惑。因为PageHelper和DataSourceHelper原理一致，都必须在执行完后立刻查询sql，如果中间出现异常，
	 * 则会出现问题。那么这个时候你就可以使用此方法。
	 * <p>同时此方法还有一点不同在于他的修改不会仅限于一次dao的查询，数据源修改的有效期会持续到整个方法结束。但是在此方法期间不允许在对数据源在进行修改，否则将会抛出异常！
	 * @param dataSourceKey
	 * @param consumer 要执行的操作
	 * @param t 方法参数
	 * @see #setDataSource(String)
	 */
	public static <T> void doQuery(String dataSourceKey,Consumer<T> consumer,T t) {
		doQuery0(dataSourceKey, (a, b) -> {
			a.accept(b);
			return null ;
		} , consumer, t);
	}
	
	/**
	 * 修改数据源并进行一次相对应的操作.
	 * <p>数据源修改相关问题可以参考{@link #setDataSource(String)}方法。
	 * <p>不同于{@link #setDataSource(String)}的是，它可以避免因PageHelper和DataSourceHelper混合使用
	 * 导致开发者不知道先执行哪个造成的困惑。因为PageHelper和DataSourceHelper原理一致，都必须在执行完后立刻查询sql，如果中间出现异常，
	 * 则会出现问题。那么这个时候你就可以使用此方法。
	 * <p>同时此方法还有一点不同在于他的修改不会仅限于一次dao的查询，数据源修改的有效期会持续到整个方法结束。但是在此方法期间不允许在对数据源在进行修改，否则将会抛出异常！
	 * @param dataSourceKey
	 * @param consumer 要执行的操作
	 * @param t 方法参数1
	 * @param u 方法参数2
	 * @see #setDataSource(String)
	 */
	public static <T,U> void doQuery(String dataSourceKey,BiConsumer<T, U> consumer,T t,U u) {
		doQuery0(dataSourceKey, (a, b) -> {
			consumer.accept(a,b);
			return null ;
		}, t, u);
	}
	
	/**
	 * 修改数据源并进行一次相对应的操作.
	 * <p>数据源修改相关问题可以参考{@link #setDataSource(String)}方法。
	 * <p>不同于{@link #setDataSource(String)}的是，它可以避免因PageHelper和DataSourceHelper混合使用
	 * 导致开发者不知道先执行哪个造成的困惑。因为PageHelper和DataSourceHelper原理一致，都必须在执行完后立刻查询sql，如果中间出现异常，
	 * 则会出现问题。那么这个时候你就可以使用此方法。
	 * <p>同时此方法还有一点不同在于他的修改不会仅限于一次dao的查询，数据源修改的有效期会持续到整个方法结束。但是在此方法期间不允许在对数据源在进行修改，否则将会抛出异常！
	 * @param dataSourceKey
	 * @param supplier 要执行的操作
	 * @return 返回指定方法执行完后的返回结果
	 * @see #setDataSource(String)
	 */
	public static <R> R doQuery(String dataSourceKey,Supplier<R> supplier) {
		return doQuery0(dataSourceKey, (t, u) -> supplier.get() , null, null);
	}
	
	/**
	 * 修改数据源并进行一次相对应的操作.
	 * <p>数据源修改相关问题可以参考{@link #setDataSource(String)}方法。
	 * <p>不同于{@link #setDataSource(String)}的是，它可以避免因PageHelper和DataSourceHelper混合使用
	 * 导致开发者不知道先执行哪个造成的困惑。因为PageHelper和DataSourceHelper原理一致，都必须在执行完后立刻查询sql，如果中间出现异常，
	 * 则会出现问题。那么这个时候你就可以使用此方法。
	 * <p>同时此方法还有一点不同在于他的修改不会仅限于一次dao的查询，数据源修改的有效期会持续到整个方法结束。但是在此方法期间不允许在对数据源在进行修改，否则将会抛出异常！
	 * @param dataSourceKey
	 * @param function 要执行的操作
	 * @param t 方法参数
	 * @return 返回指定方法执行完后的返回结果
	 * @see #setDataSource(String)
	 */
	public static <T,R> R doQuery(String dataSourceKey,Function<T, R> function,T t) {
		return doQuery0(dataSourceKey, (a, b) -> function.apply(a) , t, null);
	}
	
	/**
	 * 修改数据源并进行一次相对应的操作.
	 * <p>数据源修改相关问题可以参考{@link #setDataSource(String)}方法。
	 * <p>不同于{@link #setDataSource(String)}的是，它可以避免因PageHelper和DataSourceHelper混合使用
	 * 导致开发者不知道先执行哪个造成的困惑。因为PageHelper和DataSourceHelper原理一致，都必须在执行完后立刻查询sql，如果中间出现异常，
	 * 则会出现问题。那么这个时候你就可以使用此方法。
	 * <p>同时此方法还有一点不同在于他的修改不会仅限于一次dao的查询，数据源修改的有效期会持续到整个方法结束。但是在此方法期间不允许在对数据源在进行修改，否则将会抛出异常！
	 * @param dataSourceKey
	 * @param function 要执行的操作
	 * @param t 方法参数1
	 * @param u 方法参数2
	 * @return 返回指定方法执行完后的返回结果
	 * @see #setDataSource(String)
	 */
	public static <T,U,R> R doQuery(String dataSourceKey,BiFunction<T, U, R> function,T t,U u) {
		return doQuery0(dataSourceKey, function, t, u);
	}
	
	
	/**
	 * {@link #setDataSource(String)}方法的内部方法
	 */
	private boolean setDataSource0(String dataSourceKey,boolean inQueryCanNotChange) {
		ConnectionCache conn = this.currentDataSource.get();
		boolean isNew = conn == null;
		
		if (!this.dataSource.contains(dataSourceKey)) {
			throw new DataSourceException("指定数据源不存在：" + dataSourceKey);
		}

		if (isNew) {
			conn = new ConnectionCache(inQueryCanNotChange);
			this.currentDataSource.set(conn);
		} else {
			if (conn.isInQueryCanNotChange()) {
				throw new IllegalStateException("当前环境不允许在对数据源进行修改，可能是因为你正在执行DataSourceHelper.doQuery等相关操作。当前数据源key："
						+ conn.getKey() + "，要修改的数据源key：" + dataSourceKey);
			}
		}
		conn.setKey(dataSourceKey);
		
		if (logger.isDebugEnabled()) {
			logger.debug("修改当前操作数据源为，key={}",dataSourceKey);
		}
		
		return isNew ;
	}
	
	
	@Intercepts({ @Signature(type = StatementHandler.class, method = "prepare", args = { Connection.class, Integer.class }),
		@Signature(type = StatementHandler.class, method = "getBoundSql", args = {}), })
	class DataSourceHelperInterceptor implements Interceptor {

		/**
		 * 拦截操作只做一件事情，就是修改参数中的Connection对象.<br>
		 * 同时为了应对mybatis的执行缓存，如果修改了数据源，则会给sql语句上加注释标记
		 */
		@Override
		public Object intercept(Invocation invocation) throws Throwable {
			ConnectionCache currentConn = currentDataSource.get();
			if (currentConn == null) {
				return invocation.proceed();
			} else {
				String methodName = invocation.getMethod().getName();
				Object[] args = invocation.getArgs();
				if (methodName.equals("prepare")) {
					Connection newConn = currentConn.getConn();
					if (newConn != null) {
						if (logger.isInfoEnabled()) {
							logger.info("当前数据库操作使用的数据源被修改，key={}",currentConn.getKey());
						}
						boolean setArgs = true ;
						Connection sourceConn = (Connection) args[0];
						boolean isProxy = Proxy.isProxyClass(args[0].getClass());
						Object inv = null ;
						if (isProxy) {
							inv = Proxy.getInvocationHandler(args[0]);
							if (inv instanceof ConnectionLogger) {
								setArgs = false ;
							}
						}
						
						if (setArgs) {
							args[0] = newConn ;
							return invocation.proceed();
						} else {
							connectionField.set(inv, newConn);
							try {
								return invocation.proceed();
							} finally {
								connectionField.set(inv, sourceConn);
							}
						}
					} else {
						return invocation.proceed();
					}
				} else {
					// 标记sql语句禁止不同dataSourceKey的sql缓存(相同dataSourceKey的既然可以缓存)
					String key = currentConn.getKey();
					BoundSql sql = (BoundSql) invocation.proceed();
					if (key != null) {
						this.signBoundSql(sql, key);
					}
					return sql;
				}
			}
		}
		
		/**
		 * 在使用resumeExecutor的时候，会对Statement进行缓存，通过多sql语句进行修改，可以阻止这种缓存.
		 * 执行此方法可以将一个BoundSql中的sql语句上增加一个不影响其功能的标记，这个标记与数据源有关。
		 * 
		 * @param sql 要进行标记的BoundSql类对象
		 * @param key 数据源key
		 */
		private void signBoundSql(BoundSql sql, String key) throws IllegalAccessException {
			String sign = " /* " + key + " */ ";
			sqlField.set(sql, sql.getSql() + sign);
		}
		
		/**
		 * 只拦截StatementHandler
		 */
		@Override
		public Object plugin(Object target) {
			if (target instanceof StatementHandler) {
				return Plugin.wrap(target, this);
			} else {
				return target;
			}
		}
		
	}
	

	/**
	 * 在任何一个事务中执行此方法后，当前事务无论如何都会执行回滚操作.<br>
	 * 执行此方法的好处就是不需要抛出异常减少性能损耗
	 */
	public static void setRollbackOnly() {
		TransactionAspectSupport.currentTransactionStatus().setRollbackOnly();
	}
	

	@Override
	public void afterPropertiesSet() throws Exception {
		if (DataSourceHelper.instance != null) {
			throw new BeanInitializationException("DataSourceHelper应当是单例的，但是却被创建了两次");
		} else {
			logger.info("====> 已启用 DataSourceHelper 多数据源mybatis插件，可以使用本类实现多数据源切换操作（仅查询）");
		}
		DataSourceHelper.instance = this;
	}
	
}
